# ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv
# ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv
import warnings
warnings.filterwarnings('ignore')
import util_for_output_zh

import os
import pandas as pd
import numpy as np
from util_set_zh_matplot import plt
import seaborn as sns
from imblearn.over_sampling import SMOTE
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.ensemble import RandomForestClassifier
import xgboost as xgb
import joblib

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)

# 创建输出目录
out_dir = 'ques4_prepare_data_optim2'
os.makedirs(out_dir, exist_ok=True)


def main():
    # 准备数据（如果未生成合并文件则运行）
    prepare_data()
    # if not os.path.exists(os.path.join(out_dir, 'ques4_concat.csv')):
    #     prepare_data()
    # stage1()


def prepare_data():
    """数据预处理与合并"""
    # 加载数据（使用正斜杠路径）
    female_data = pd.read_csv('ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv')
    male_data = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

    # 去除未命名列
    def remove_unnamed_columns(df):
        unnamed_cols = [col for col in df.columns if 'Unnamed' in col]
        return df.drop(columns=unnamed_cols, errors='ignore')
    
    female_data = remove_unnamed_columns(female_data)
    male_data = remove_unnamed_columns(male_data)

    # 统一列名
    male_data = male_data.rename(columns={'唯一比对的读段数  ': '唯一比对的读段数'})
    
    # 处理性别特有列
    female_specific_cols = ['孕周天数']
    male_specific_cols = ['实际孕周', '达标孕周', '达标', 'Y染色体的Z值', 'Y染色体浓度']
    female_data = female_data.drop(columns=female_specific_cols, errors='ignore')
    male_data = male_data.drop(columns=male_specific_cols, errors='ignore')
    
    # 添加性别列并合并
    female_data['性别'] = '女'
    male_data['性别'] = '男'
    combined_data = pd.concat([female_data, male_data], axis=0, ignore_index=True)
    
    # 保存合并数据
    combined_data.to_csv(os.path.join(out_dir, 'ques4_concat.csv'), 
                         index=False, encoding='utf-8-sig')
    print("数据已合并并保存为 ques4_concat.csv")
    print(f"总记录数: {len(combined_data)} (女胎: {len(female_data)}, 男胎: {len(male_data)})")

    # 验证列名一致性
    if set(female_data.columns) == set(male_data.columns):
        print("列名验证通过：两个数据集列名一致")
    else:
        diff_cols = set(female_data.columns) ^ set(male_data.columns)
        print(f"警告：列名存在差异 {diff_cols}")

    # 绘制类别分布图表
    plot_category_distributions(female_data, male_data)


def plot_category_distributions(female_data, male_data):
    """绘制男女及总体类别分布"""
    female_counts = female_data['染色体的非整倍体'].value_counts()
    male_counts = male_data['染色体的非整倍体'].value_counts()
    total_counts = female_counts.add(male_counts, fill_value=0).sort_values(ascending=False)

    plt.figure(figsize=(15, 12))
    
    plt.subplot(3, 1, 1)
    female_counts.plot(kind='bar', color='pink')
    plt.title('女胎染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.subplot(3, 1, 2)
    male_counts.plot(kind='bar', color='lightblue')
    plt.title('男胎染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.subplot(3, 1, 3)
    total_counts.plot(kind='bar', color='purple')
    plt.title('总体染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '类别分布统计.png'))
    plt.close()
    print("类别分布图表已保存")


def cross_gender_test(balanced_data, features, le, train_gender, test_gender):
    """跨性别预测通用函数"""
    # 划分训练集和测试集
    train_mask = balanced_data['性别'] == train_gender
    test_mask = balanced_data['性别'] == test_gender
    
    X_train = balanced_data.loc[train_mask, features]
    y_train = le.transform(balanced_data.loc[train_mask, '染色体的非整倍体'])
    X_test = balanced_data.loc[test_mask, features]
    y_test = le.transform(balanced_data.loc[test_mask, '染色体的非整倍体'])

    # 数据标准化（仅用训练集拟合）
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 使用XGBoost模型（性能更优）
    clf = xgb.XGBClassifier(
        objective='multi:softmax',
        num_class=len(le.classes_),
        n_estimators=150,
        max_depth=5,
        learning_rate=0.1,
        scale_pos_weight=1,
        random_state=42
    )
    clf.fit(X_train_scaled, y_train)
    y_pred = clf.predict(X_test_scaled)

    # 输出评估结果
    print(f"\n{train_gender}训练集预测{test_gender}测试集结果:")
    print(classification_report(y_test, y_pred, target_names=le.classes_))
    
    # 绘制混淆矩阵
    plot_confusion_matrix(y_test, y_pred, le.classes_, 
                         f"{train_gender}→{test_gender}混淆矩阵")
    
    return clf


def plot_confusion_matrix(y_true, y_pred, labels, title):
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=labels, yticklabels=labels)
    plt.xlabel('预测类别')
    plt.ylabel('实际类别')
    plt.title(title)
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, f'{title}.png'))
    plt.close()


def stage1():
    """模型训练与评估主流程"""
    # 加载数据（修复路径转义问题）
    file_path = 'ques4_prepare_data/ques4_concat.csv'
    combined_data = pd.read_csv(file_path)

    # 分离男女数据
    female_data = combined_data[combined_data['性别'] == '女'].copy()
    male_data = combined_data[combined_data['性别'] == '男'].copy()

    # 过滤有效类别（样本数≥10且男女均存在）
    total_counts = combined_data['染色体的非整倍体'].value_counts()
    valid_categories = [label for label in total_counts.index 
                       if total_counts[label] >= 10 
                       and label in female_data['染色体的非整倍体'].unique()
                       and label in male_data['染色体的非整倍体'].unique()]
    
    print(f"原始类别分布:\n{total_counts}\n")
    print(f"过滤后保留的类别:\n{valid_categories}\n")

    # 筛选数据
    final_data = combined_data[combined_data['染色体的非整倍体'].isin(valid_categories)]
    print(f"最终使用的共同类别:\n{valid_categories}\n")

    # 定义特征列
    features = [
        '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', '唯一比对的读段数',
        'GC含量', '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
        'X染色体的Z值', 'X染色体浓度', '13号染色体的GC含量', 
        '18号染色体的GC含量', '21号染色体的GC含量'
    ]

    # 检查特征存在性
    missing_features = [f for f in features if f not in final_data.columns]
    if missing_features:
        raise ValueError(f"特征不存在: {missing_features}")

    # 平衡数据（使用SMOTE过采样而非复制）
    balanced_data = pd.DataFrame()
    max_samples = 700  # 每个类别最大样本数
    
    for gender in ['女', '男']:
        gender_data = final_data[final_data['性别'] == gender].copy()
        # 按类别采样到最小样本数的2倍（为SMOTE留空间）
        min_samples = gender_data['染色体的非整倍体'].value_counts().min()
        sample_size = min(min_samples * 2, max_samples)
        
        # 下采样到sample_size
        gender_downsampled = pd.DataFrame()
        for label in valid_categories:
            label_subset = gender_data[gender_data['染色体的非整倍体'] == label]
            take = min(len(label_subset), sample_size)
            gender_downsampled = pd.concat([
                gender_downsampled,
                label_subset.sample(take, random_state=42)
            ], ignore_index=True)
        
        # 对下采样后的数据用SMOTE过采样到目标数量
        X = gender_downsampled[features]
        y = gender_downsampled['染色体的非整倍体']
        le_temp = LabelEncoder()
        y_encoded = le_temp.fit_transform(y)
        
        # 计算需要生成的样本数
        n_total = len(valid_categories) * max_samples
        smote = SMOTE(sampling_strategy='auto', random_state=42)
        X_resampled, y_resampled = smote.fit_resample(X, y_encoded)
        
        # 截断到目标数量
        resampled = pd.DataFrame(X_resampled, columns=features)
        resampled['染色体的非整倍体'] = le_temp.inverse_transform(y_resampled)
        resampled['性别'] = gender
        resampled = resampled.groupby('染色体的非整倍体').apply(
            lambda x: x.sample(min(len(x), max_samples), random_state=42)
        ).reset_index(drop=True)
        
        balanced_data = pd.concat([balanced_data, resampled], ignore_index=True)
    # 打印平衡前后分布
    print("平衡前类别分布:")
    print(final_data['染色体的非整倍体'].value_counts())
    print("\n平衡后类别分布:")
    print(balanced_data['染色体的非整倍体'].value_counts())

    # 标签编码
    le = LabelEncoder()
    le.fit(valid_categories)  # 确保编码一致性

    # 跨性别测试
    clf_male2female = cross_gender_test(balanced_data, features, le, '男', '女')
    clf_female2male = cross_gender_test(balanced_data, features, le, '女', '男')

    # 特征重要性分析（使用最后一个训练的模型）
    print("\n特征重要性排序:")
    importance = pd.DataFrame({
        'feature': features,
        'importance': clf_female2male.feature_importances_
    }).sort_values('importance', ascending=False)
    print(importance)

    # 保存模型
    joblib.dump(clf_male2female, os.path.join(out_dir, 'male2female_model.pkl'))
    joblib.dump(clf_female2male, os.path.join(out_dir, 'female2male_model.pkl'))
    print("\n模型已保存至 ques4_prepare_data 目录")


if __name__ == '__main__':
    main()